from typing import Annotated, Sequence, Literal
from typing_extensions import TypedDict
from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph.message import add_messages
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from customize.get_ollama import GetOllama
from pydantic import BaseModel, Field
from langgraph.prebuilt import tools_condition

from langchain.tools.retriever import create_retriever_tool
import pprint
from langgraph.graph import END, StateGraph, START
from langgraph.prebuilt import ToolNode
from langchain_chroma import Chroma
from customize.langchain_customized_huggingface import HuggingFaceBgeEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

model_name = "F:/models/BAAI/bge-large-zh-v1.5"
# model_name = "BAAI/bge-large-zh-v1.5"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': True}  # set True to compute cosine similarity

embeddings = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    model_path=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
    query_instruction="为这个句子生成表示以用于检索相关文章："
)

from langchain_community.document_loaders import TextLoader
loader = TextLoader('resources/state_of_the_union.txt', encoding='utf-8')
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=100, chunk_overlap=50
)
doc_splits = text_splitter.split_documents(documents)
persist_directory = 'db'
# Add to vectorDB
# vectorstore = Chroma.from_documents(
#     documents=doc_splits,
#     collection_name="rag-chroma",
#     embedding=embeddings,
#     persist_directory=persist_directory
# )
vectorstore = Chroma(collection_name="jpjd", embedding_function=embeddings, persist_directory=persist_directory, collection_metadata={"hnsw:batch_size": 10000})
retriever = vectorstore.as_retriever()

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]


retriever_tool = create_retriever_tool(
    retriever,
    "极品家丁检索器",
    "此工具专门用于在小说《极品家丁》中进行信息检索，可依据用户提供的关键词、情节片段、人物名称等，从小说文本中查找与之相关的段落、情节发展、人物关系等详细信息，助力用户深入探索《极品家丁》这部小说的丰富内容。小说的主要人物有林三（原名林晚荣），肖青璇等。",

)
tools = [retriever_tool]


# Nodes
def agent(state):

    print("--call agent--")
    messages = state['messages']
    model = GetOllama(model_type=1, streaming=True, temperature=0)()
    model.bind_tools(tools)
    response = model.invoke(messages)
    return {"messages": [response]}


def rewrite(state):

    print("---TRANSFORM QUERY---")
    messages = state['messages']
    question = messages[0].content

    msg = [
        HumanMessage(
            content=f""" \n 
        Look at the input and try to reason about the underlying semantic intent / meaning. \n 
        Here is the initial question:
        \n ------- \n
        {question} 
        \n ------- \n
        Formulate an improved question: """,
        )
    ]
    model = GetOllama(model_type=1, streaming=True, temperature=0)()
    response = model.invoke(msg)
    return {"messages": [response]}


def generate(state):
    print("---GENERATE---")
    messages = state["messages"]
    question = messages[0].content
    last_message = messages[-1]

    docs = last_message.content

    # Prompt
    prompt = hub.pull("rlm/rag-prompt")

    # LLM
    llm = GetOllama(model_type=1, temperature=0, streaming=True)()

    # Post-processing
    def format_docs(docs):
        return "\n\n".join(doc.page_content for doc in docs)

    # Chain
    rag_chain = prompt | llm | StrOutputParser()

    # Run
    response = rag_chain.invoke({"context": docs, "question": question})
    return {"messages": [response]}


# Edges


def grade_documents(state) -> Literal['generate', 'rewrite']:

    print("---CHECK RELEVANCE---")

    # Data model
    class grade(BaseModel):
        """Binary score for relevance check."""

        binary_score: str = Field(description="Relevance score 'yes' or 'no'")

    # LLM
    model = GetOllama(temperature=0, model_type=1, streaming=True)()

    # LLM with tool and validation
    llm_with_tool = model.with_structured_output(grade)

    # Prompt
    prompt = PromptTemplate(
        template="""You are a grader assessing relevance of a retrieved document to a user question. \n 
            Here is the retrieved document: \n\n {context} \n\n
            Here is the user question: {question} \n
            If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
            Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.""",
        input_variables=["context", "question"],
    )

    # Chain
    chain = prompt | llm_with_tool

    messages = state["messages"]
    last_message = messages[-1]

    question = messages[0].content
    docs = last_message.content

    scored_result = chain.invoke({"question": question, "context": docs})

    score = scored_result.binary_score

    if score == "yes":
        print("---DECISION: DOCS RELEVANT---")
        return "generate"

    else:
        print("---DECISION: DOCS NOT RELEVANT---")
        print(score)
        return "rewrite"


workflow = StateGraph(AgentState)

workflow.add_node("agent", agent)
retrieve = ToolNode([retriever_tool])
workflow.add_node("retrieve", retrieve)
workflow.add_node("rewrite", rewrite)
workflow.add_node("generate", generate)

workflow.add_edge(START, "agent")
workflow.add_conditional_edges(
    "agent",
    tools_condition,
    {
        "tools": "retrieve",
        END: END,
    },
)
# Edges taken after the `action` node is called.
workflow.add_conditional_edges(
    "retrieve",
    # Assess agent decision
    grade_documents,
)
workflow.add_edge("generate", END)
workflow.add_edge("rewrite", "agent")

graph = workflow.compile()

from IPython.display import Image
from PIL import Image as PILImage
import io
img = Image(graph.get_graph().draw_mermaid_png())
pil_img = PILImage.open(io.BytesIO(img.data))
pil_img.save("save_image.png")

inputs = {
    "messages": [
        # ("user", "What does Lilian Weng say about the types of agent memory?"),
        ("user", "极品家丁中，林三和青璇的关系?")
    ]
}
for output in graph.stream(inputs):
    for key, value in output.items():
        pprint.pprint(f"Output from node '{key}':")
        pprint.pprint("---")
        pprint.pprint(value, indent=2, width=80, depth=None)
    pprint.pprint("\n---\n")